use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    pub from: String,
    pub value: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Conversation {
    pub conversations: Vec<Message>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub system: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tools: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Dataset {
    pub conversations: Vec<Conversation>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationResult {
    pub is_valid: bool,
    pub errors: Vec<String>,
    pub warnings: Vec<String>,
}

impl Conversation {
    #[allow(dead_code)]
    pub fn new() -> Self {
        Self {
            conversations: vec![
                Message {
                    from: "human".to_string(),
                    value: String::new(),
                },
                Message {
                    from: "gpt".to_string(),
                    value: String::new(),
                },
            ],
            system: None,
            tools: None,
        }
    }

    pub fn validate(&self) -> ValidationResult {
        let mut errors = Vec::new();
        let mut warnings = Vec::new();

        // 验证消息数组不为空
        if self.conversations.is_empty() {
            warnings.push("对话为空".to_string());
            return ValidationResult {
                is_valid: true,
                errors,
                warnings,
            };
        }

        // 统计各角色数量
        let mut human_count = 0;
        let mut gpt_count = 0;
        let mut function_call_count = 0;
        let mut observation_count = 0;

        // 验证消息序列和角色配对
        for (i, message) in self.conversations.iter().enumerate() {
            let role = &message.from;
            
            match role.as_str() {
                "human" => human_count += 1,
                "gpt" => gpt_count += 1,
                "thinking" => {
                    // thinking 可以出现在任何位置，无特殊约束
                }
                "function_call" => {
                    function_call_count += 1;
                    // 检查是否有对应的observation
                    if i == self.conversations.len() - 1 
                        || self.conversations[i + 1].from != "observation" {
                        errors.push(format!("第{}个消息: function_call 后必须跟随 observation", i + 1));
                    }
                }
                "observation" => {
                    observation_count += 1;
                    // 检查是否有对应的function_call
                    if i == 0 || self.conversations[i - 1].from != "function_call" {
                        errors.push(format!("第{}个消息: observation 前必须有 function_call", i + 1));
                    }
                }
                _ => {
                    errors.push(format!("第{}个消息: 未知的角色类型 \"{}\"", i + 1, role));
                }
            }

            // 检查消息内容是否为空
            if message.value.trim().is_empty() {
                warnings.push(format!("第{}个消息 ({}) 内容为空", i + 1, role));
            }
        }

        // 检查human和gpt的配对关系
        if human_count != gpt_count {
            errors.push(format!(
                "human ({}) 和 gpt ({}) 消息数量必须相等",
                human_count, gpt_count
            ));
        }

        // 检查function_call和observation的配对关系
        if function_call_count != observation_count {
            errors.push(format!(
                "function_call ({}) 和 observation ({}) 消息数量必须相等",
                function_call_count, observation_count
            ));
        }

        ValidationResult {
            is_valid: errors.is_empty(),
            errors,
            warnings,
        }
    }
}

impl Dataset {
    pub fn new() -> Self {
        Self {
            conversations: Vec::new(),
        }
    }

    pub fn add_conversation(&mut self, conversation: Conversation) {
        self.conversations.push(conversation);
    }

    pub fn validate(&self) -> ValidationResult {
        let mut all_errors = Vec::new();
        let mut all_warnings = Vec::new();

        if self.conversations.is_empty() {
            all_warnings.push("数据集为空".to_string());
        }

        for (i, conversation) in self.conversations.iter().enumerate() {
            let result = conversation.validate();
            
            // 为错误信息添加对话索引前缀
            for error in result.errors {
                all_errors.push(format!("对话{}: {}", i + 1, error));
            }
            
            for warning in result.warnings {
                all_warnings.push(format!("对话{}: {}", i + 1, warning));
            }
        }

        ValidationResult {
            is_valid: all_errors.is_empty(),
            errors: all_errors,
            warnings: all_warnings,
        }
    }
} 